package com.shiyi.shiyiaicodefather.core.handle;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.io.FileUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONObject;
import cn.hutool.json.JSONUtil;
import com.shiyi.shiyiaicodefather.constant.AppConstant;
import com.shiyi.shiyiaicodefather.core.builder.VueProjectBuilder;
import com.shiyi.shiyiaicodefather.model.entity.ChatHistory;
import com.shiyi.shiyiaicodefather.model.entity.User;
import com.shiyi.shiyiaicodefather.model.enums.ChatHistoryMessageTypeEnum;
import com.shiyi.shiyiaicodefather.model.message.*;
import com.shiyi.shiyiaicodefather.service.ChatHistoryService;
import com.shiyi.shiyiaicodefather.tool.BaseTool;
import com.shiyi.shiyiaicodefather.tool.ToolManager;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.io.File;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

/**
 * JSON 消息流处理器
 * 处理 VUE_PROJECT 类型的复杂流式响应，包含工具调用信息
 */
@Slf4j
@Component
public class JsonMessageStreamHandler {

  @Resource
  private VueProjectBuilder vueProjectBuilder;

  @Resource
  private ToolManager toolManager;


    /**
     * 处理 TokenStream（VUE_PROJECT）
     * 解析 JSON 消息并重组为完整的响应格式
     *
     * @param originFlux         原始流
     * @param chatHistoryService 聊天历史服务
     * @param appId              应用ID
     * @param loginUser          登录用户
     * @return 处理后的流
     */
    public Flux<String> handle(Flux<String> originFlux,
                               ChatHistoryService chatHistoryService,
                               long appId, User loginUser) {
        // 收集数据用于生成后端记忆格式
        StringBuilder chatHistoryToDisplay = new StringBuilder();

        // 收集数据用于生成后端记忆格式
        StringBuilder chatHistoryToCache = new StringBuilder();
        // 每个 Flux 流可能包含多条工具调用和 AI_RESPONSE 响应信息，统一收集之后批量入库
        List<ChatHistory> chatHistoryList = new ArrayList<>();
        // 用于跟踪已经见过的工具ID，判断是否是第一次调用
        Set<String> seenToolIds = new HashSet<>();
        return originFlux
                .map(chunk -> {
                    // 解析每个 JSON 消息块
                    return handleJsonMessageChunk(chunk, chatHistoryToDisplay,chatHistoryToCache,chatHistoryList,seenToolIds);
                })
                .filter(StrUtil::isNotEmpty) // 过滤空字串
                .doOnComplete(() -> {
                  //流式响应完成后

                  //添加工具相关信息到对话历史
                  if (CollUtil.isNotEmpty(chatHistoryList)) {
                    chatHistoryList.forEach(item -> {
                      item.setAppId(appId);
                      item.setUserId(loginUser.getId());
                    });

                    chatHistoryService.addChatMessageBatch(chatHistoryList);
                  }else{
                    //如果没有工具调用
                    String cacheString = chatHistoryToCache.toString();
                    chatHistoryService.addChatMessage(appId, cacheString, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                  }


                    // 添加 AI 消息到对话历史
                    String aiResponse = chatHistoryToDisplay.toString();
                    chatHistoryService.addChatMessage(appId, aiResponse, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                    String projectPath = AppConstant.CODE_OUTPUT_ROOT_DIR + File.separator + "vue_project_" + appId;
                    vueProjectBuilder.buildProjectAsync(projectPath);
                })
                .doOnError(error -> {
                    // 如果AI回复失败，也要记录错误消息
                    String errorMessage = "AI回复失败: " + error.getMessage();
                    chatHistoryService.addChatMessage(appId, errorMessage, ChatHistoryMessageTypeEnum.AI.getValue(), loginUser.getId());
                });
    }

    /**
     * 解析并收集 TokenStream 数据
     */
    private String handleJsonMessageChunk(String chunk, StringBuilder chatHistoryToDisplay,StringBuilder chatHistoryToCache,List<ChatHistory> chatHistoryList,Set<String> seenToolIds) {
        // 解析 JSON
        StreamMessage streamMessage = JSONUtil.toBean(chunk, StreamMessage.class);
        StreamMessageTypeEnum typeEnum = StreamMessageTypeEnum.getEnumByValue(streamMessage.getType());
        switch (typeEnum) {
            case AI_RESPONSE -> {
                AiResponseMessage aiMessage = JSONUtil.toBean(chunk, AiResponseMessage.class);
                String data = aiMessage.getData();
                // 直接拼接响应
              chatHistoryToDisplay.append(data);
              chatHistoryToCache.append(data);
              return data;
            }
            case TOOL_REQUEST -> {
                ToolRequestMessage toolRequestMessage = JSONUtil.toBean(chunk, ToolRequestMessage.class);
                String toolId = toolRequestMessage.getId();
                String toolName = toolRequestMessage.getName();
              // 检查是否是第一次看到这个工具 ID
                if (toolId != null && !seenToolIds.contains(toolId)) {
                    // 第一次调用这个工具，记录 ID 并完整返回工具信息
                    seenToolIds.add(toolId);
                  BaseTool tool = toolManager.getTool(toolName);
                  return tool.generateToolRequestResponse();
                } else {
                    // 不是第一次调用这个工具，直接返回空
                    return "";
                }
            }
            case TOOL_EXECUTED -> {
              processToolExecutionMessage(chunk, chatHistoryToCache, chatHistoryList);

              ToolExecutedMessage toolExecutedMessage = JSONUtil.toBean(chunk, ToolExecutedMessage.class);
                JSONObject jsonObject = JSONUtil.parseObj(toolExecutedMessage.getArguments());
                String toolName = toolExecutedMessage.getName();
                BaseTool tool = toolManager.getTool(toolName);
                String result = tool.generateToolExecutedResult(jsonObject);
                // 输出前端和要持久化的内容
                String output = String.format("\n\n%s\n\n", result);
                chatHistoryToDisplay.append(output);
                return output;
            }
            default -> {
                log.error("不支持的消息类型: {}", typeEnum);
                return "";
            }
        }
    }

    /**
     * 解析处理工具调用相关信息
     * @param chunk
     * @param chatHistoryToCache
     * @param chatHistoryList
     */
    private static void processToolExecutionMessage(String chunk, StringBuilder chatHistoryToCache, List<ChatHistory> chatHistoryList) {
      ToolExecutedMessage toolExecutedMessage1 = JSONUtil.toBean(chunk, ToolExecutedMessage.class);
      String cacheString = chatHistoryToCache.toString();
      ToolRequestMessage toolRequestMessage = new ToolRequestMessage();
      toolRequestMessage.setId(toolExecutedMessage1.getId());
      toolRequestMessage.setName(toolExecutedMessage1.getName());
      toolRequestMessage.setArguments(toolExecutedMessage1.getArguments());
      toolRequestMessage.setText(cacheString);
      String jsonStr = JSONUtil.toJsonStr(toolRequestMessage);
      ChatHistory toolRequestHistory = ChatHistory.builder()
        .message(jsonStr)
        .messageType(ChatHistoryMessageTypeEnum.TOOL_EXECUTION_REQUEST.getValue())
        .category("cache")
        .build();
      chatHistoryList.add(toolRequestHistory);
      ChatHistory toolResultHistory = ChatHistory.builder()
        .message(chunk)
        .messageType(ChatHistoryMessageTypeEnum.TOOL_EXECUTION_RESULT.getValue())
        .category("cache")
        .build();
      chatHistoryList.add(toolResultHistory);
      chatHistoryToCache.setLength(0);
    }
}
